import platform
import shutil
from datetime import datetime
from multiprocessing import freeze_support, cpu_count, Pool

import pandas as pd
import akshare as ak
import os
# from stock_tools.Config import update_daily_dir, data_from_dir

update_daily_dir = r'D:/python/xbx_stock_2021_part3/市值选谷策略深度优化实盘/data/daily_data/xbx增量日线/'
data_from_dir = r'D:/python/xbx_stock_2021_part3/市值选谷策略深度优化实盘/data/daily_data/stock/'


def update_single_stock_daily(file_name: str):
    df_new = pd.read_csv(update_daily_dir + file_name, encoding='gbk', skiprows=1, engine='pyarrow')
    df_new['交易日期'] = pd.to_datetime(df_new['交易日期'])
    if file_name in os.listdir(data_from_dir):
        df_old = pd.read_csv(data_from_dir + file_name, encoding='gbk')
        df_old['交易日期'] = pd.to_datetime(df_old['交易日期'])
        # 虽然原定的是一周数据, 但是还是用条件筛选一下
        # 获取旧数据最后一天的日期
        last_day = df_old.iloc[-1]['交易日期']
        # 获取大于这个日期的新数据
        df_new = df_new[df_new['交易日期'] > pd.to_datetime(last_day)]
        # 获取好了数据之后直接concat
        df_old = pd.concat([df_old, df_new])
        # 输出替换原先的数据文件
        df_old.to_csv(data_from_dir + file_name, encoding='gbk', index=False)
        print(f"更新 {file_name} 股票数据完成!")
    else:
        # 如果不在之前的列表说明是新股，直接保存就完事
        df_new.to_csv(data_from_dir + file_name, encoding='gbk', index=False)
        print(f"添加 {file_name} 股票数据完成!")


def get_index_stock_data(out_data_dir_path):
    # 第一步 获取实时行情数据 拿到所有的指数代码
    df = ak.stock_zh_index_spot()

    index_dict = {
        'sz980015': '疫苗生科',
        'sz980016': '公卫健康',
        'sz980028': '龙头家电',
        'sz980030': '消费电子',
        'sz980032': '新能电池'
    }

    for index in range(df.shape[0]):
        index_symbol = df.loc[index, '代码']
        index_name = df.loc[index, '名称']
        index_df = ak.stock_zh_index_daily(index_symbol)
        index_df['指数代码'] = index_symbol
        index_df['指数名称'] = index_name if index_name != "" else index_dict[index_symbol]
        rename_dict = {'date': '交易日期', 'open': '开盘价', 'high': '最高价', 'low': '最低价', 'close': '收盘价',
                       'volume': '成交量'}
        # 其中amount单位是手，说明数据不够精确
        index_df.rename(columns=rename_dict, inplace=True)
        index_df['前收盘价'] = index_df['收盘价'].shift()
        index_df.loc[0, '前收盘价'] = index_df.loc[0, '开盘价']

        index_df = index_df[
            ['指数代码', '交易日期', '指数名称', '开盘价', '最高价', '最低价', '收盘价', '成交量', '前收盘价']]
        index_df.to_csv(out_data_dir_path + index_symbol + '.csv', encoding='gbk', index=False)
        print(f"{index_symbol} - {index_name} 数据导出成功!")
        time.sleep(0.2)


if __name__ == '__main__':
    # 模式1为更新所有日线数据 模式2为更新所有指数日线数据
    MODE = 1
    if MODE == 1:
        print("更新所有股票日线数据中...")
        # 更新股票日线数据
        stock_file_list = os.listdir(update_daily_dir)
        if 'Windows' in platform.platform():
            freeze_support()
        # 标记开始时间
        start_time = datetime.now()

        # ===并行提速的办法
        with Pool(max(cpu_count() - 4, 1)) as pool:
            df_list = pool.map(update_single_stock_daily, sorted(stock_file_list))
        print('更新日线数据完毕，消耗时间', datetime.now() - start_time)
    elif MODE == 2:
        print("更新所有指数日线数据中...")
        out_index_stock_data_dir_path = os.path.join('D:/python/xbx_stock_2021_part3/市值选谷策略深度优化实盘/data/index_data')
        # 删除所有指数数据
        shutil.rmtree(out_index_stock_data_dir_path)
        os.mkdir(out_index_stock_data_dir_path)
        # 获取所有指数数据
        get_index_stock_data(out_index_stock_data_dir_path)
    else:
        print("未知的命令")
